"""
数据库连接管理
"""
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.pool import QueuePool
from config.settings import settings
from config.database import Base
from typing import AsyncGenerator
import logging
from sqlalchemy import text

logger = logging.getLogger(__name__)


class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self):
        self.engine = None
        self.SessionLocal = None
        self._initialize()
    
    def _initialize(self):
        """初始化数据库连接"""
        try:
            # 创建异步数据库引擎
            self.engine = create_async_engine(
                settings.database_url,
                echo=settings.debug,
                poolclass=QueuePool,
                pool_size=20,  # 增加基础连接池大小
                max_overflow=30,  # 增加溢出连接数
                pool_pre_ping=True,
                pool_recycle=1800,  # 30分钟回收连接，缩短回收时间
                pool_timeout=60,  # 增加超时时间
                connect_args={
                    "charset": "utf8mb4",
                    "autocommit": False
                }
            )
            
            # 创建异步会话工厂
            self.SessionLocal = async_sessionmaker(
                autocommit=False,
                autoflush=False,
                bind=self.engine,
                class_=AsyncSession
            )
            
            logger.info("异步数据库连接初始化成功")
            
        except Exception as e:
            logger.error(f"异步数据库连接初始化失败: {str(e)}")
            raise
    
   
    
    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        """获取异步数据库会话"""
        session = None
        try:
            session = self.SessionLocal()
            yield session
        except Exception as e:
            if session:
                try:
                    await session.rollback()
                except Exception as rollback_e:
                    logger.error(f"回滚失败: {str(rollback_e)}")
            logger.error(f"数据库会话错误: {str(e)}")
            raise
        finally:
            if session:
                try:
                    await session.close()
                except Exception as close_e:
                    logger.error(f"关闭会话失败: {str(close_e)}")
    
    async def close_all_connections(self):
        """关闭所有数据库连接"""
        if self.engine:
            await self.engine.dispose()
            logger.info("所有数据库连接已关闭")


# 全局数据库管理器实例
db_manager = DatabaseManager()


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """获取异步数据库会话的依赖注入函数"""
    async for session in db_manager.get_session():
        yield session
